import wave
import os
from tqdm import tqdm
import sys
import re
import codecs
import psutil
import multiprocessing
from multiprocessing import Pool, Queue, Process


def is_wav(f):
    res = True
    try:
        wave.open(f)
    except wave.Error as e:
        res = False
    return res


def pcm2wav(pcm_file, save_file, channels=1, bits=16, sample_rate=16000):
    """ pcm转换为wav格式
        Args:
            pcm_file pcm文件
            save_file 保存文件
            channels 通道数
            bits 量化位数，即每个采样点占用的比特数
            sample_rate 采样频率
    """
    if is_wav(pcm_file):
        raise ValueError('"' + str(pcm_file) + '"' +
                         " is a wav file, not pcm file! ")

    pcmf = open(pcm_file, 'rb')
    pcmdata = pcmf.read()
    pcmf.close()

    if bits % 8 != 0:
        raise ValueError("bits % 8 must == 0. now bits:" + str(bits))

    wavfile = wave.open(save_file, 'wb')

    wavfile.setnchannels(channels)
    wavfile.setsampwidth(bits // 8)
    wavfile.setframerate(sample_rate)

    wavfile.writeframes(pcmdata)
    wavfile.close()


def convert_dir(root, ext=".pcm", **kwargs):
    # convert_dir(r'/path/to/pcm/dir', '.pcm')
    src_files = [os.path.join(dir_path, f)
                 for dir_path, _, files in os.walk(root)
                 for f in files
                 if os.path.splitext(f)[1] == ext]

    for src_file in tqdm(src_files, ascii=True):
        try:
            wav_file = os.path.splitext(src_file)[0] + ".wav"
            pcm2wav(src_file, wav_file, **kwargs)
        except Exception as e:
            print('Convert fail: ' + src_file)
            print(e)


def main(zip_file_list):
    with codecs.open(zip_file_list, 'r') as f:
        for src_file in f:
            sources_dir = src_file.strip().replace('.zip', '')
            print('Convert source: ' + sources_dir)
            convert_dir(sources_dir, '.pcm')


def conv_worker(sub_dirs_list, thread_id):
    print("Converting worker threads: {}, tasks: {}".format(
        thread_id, len(sub_dirs_list)))
    for src_file in sub_dirs_list:
        try:
            convert_dir(src_file, '.pcm')
        except Exception as e:
            print('Convert fail: ' + src_file)
            continue


def multiply_workers(zip_file_list):

    cpu_number = psutil.cpu_count()
    dirs_list = []
    with codecs.open(zip_file_list, 'r') as f:
        for src_file in f:
            sources_dir = src_file.strip().replace('.zip', '')
            print('Convert source: ' + sources_dir)
            # convert_dir(sources_dir, '.pcm')
            dirs_list.append(sources_dir)
    process_list = []
    tasks_every_worker = len(dirs_list) // cpu_number + 1
    for i in range(cpu_number):  # 开启5个子进程执行fun1函数
        sub_tasks = dirs_list[i * tasks_every_worker: min(
            tasks_every_worker * (i + 1), len(dirs_list))]
        p = Process(target=conv_worker, args=(sub_tasks, i, ))  # 实例化进程对象
        p.start()
        process_list.append(p)

    for i in process_list:
        p.join()


if __name__ == '__main__':
    multiply_workers(sys.argv[1])
